# -*- coding: utf-8 -*-
"""
可用性节点替换
1.找出http/ping监测，area/area_ipv6存在旧节点的asset_tasks
2.处理asset_tasks的area/area_ipv6旧节点为新节点
3.被修改且监测开启的asset_tasks修改配置下发至scanner
"""
from flask import current_app as app
from mongoengine.queryset.visitor import Q
from app.db.models.asset_tasks import AssetTasks
from app.db.models.wass import ScanEvents
from app.libs.utils import enable_available_task, enable_web_sec_task
from app.libs.utility import IPV6_AREA_LIST


class AvailableNode:

    def __init__(self, old, new):
        self.old_node = old
        self.new_node = new

    @staticmethod
    def log_process(count, total):
        count += 1
        app.logger.info(f"process {count}/{total}")
        return count

    def log(self):
        self.get_asset_task_ids()

    def get_asset_task_ids(self):
        asset_task_ids = []
        for task_type in ["http", "ping"]:
            query = Q(**{
                "taskGroupType": task_type,
                f"{task_type}Settings__collect__area": self.old_node
            }) | Q(**{
                "taskGroupType": task_type,
                f"{task_type}Settings__collect__area_ipv6": self.old_node
            })
            ids = AssetTasks.objects.filter(query).distinct("_id")
            asset_task_ids += ids
            app.logger.info(f"{task_type}监测 {len(ids)}个asset_tasks需处理")
        return asset_task_ids

    def run(self, enable=False):
        asset_task_ids = self.get_asset_task_ids()
        all_count = len(asset_task_ids)
        count = 0
        update_enable_asset_task_ids = set()
        for _ in asset_task_ids:
            count = self.log_process(count, all_count)
            asset_task = AssetTasks.objects.filter(id=_).first()
            settings = getattr(asset_task, f"{asset_task['taskGroupType']}Settings")
            for area in [settings.collect.area, settings.collect.area_ipv6]:
                if area and self.old_node in area:
                    app.logger.info(f"Old node '{self.old_node}' in area {area}")
                    index = area.index(self.old_node)
                    if self.new_node not in area:
                        area[index] = self.new_node
                        app.logger.info(f"Replace old node '{self.old_node}' to new node {self.new_node}")
                    else:
                        for area_node in IPV6_AREA_LIST:
                            if area_node not in area:
                                area[index] = area_node
                                app.logger.info(f"Replace old node '{self.old_node}' to new node '{area_node}'")
                                break
                        if self.old_node in area:
                            area.remove(self.old_node)
                            app.logger.info(f"No new node can replace, remove old node '{self.old_node}'")
                    app.logger.info(f"New config area is {area}")
            asset_task.save()
            if asset_task.enableMonitor and settings.enable and asset_task.triggerType == "schedule":
                app.logger.info(f"{asset_task.id} need update scanner")
                update_enable_asset_task_ids.add(asset_task.id)
        if enable and update_enable_asset_task_ids:
            self.enable_tasks(list(update_enable_asset_task_ids))
            return set()
        return update_enable_asset_task_ids

    @staticmethod
    def enable_tasks(update_enable_asset_task_ids: list):
        asset_tasks = AssetTasks.objects.filter(
            id__in=update_enable_asset_task_ids,
            taskGroupType__in=['http', 'ping'],
            enableMonitor=True
        )
        total = asset_tasks.count()
        count = 0
        for asset_task in asset_tasks:
            count = AvailableNode.log_process(count, total)
            settings = getattr(asset_task, f"{asset_task.taskGroupType}Settings")
            try:
                enable_available_task(asset_task)
            except Exception as e:
                app.logger.warning(f"{asset_task.taskGroupType} asset_task: {asset_task.id} "
                                   f"task_id: {settings.taskId} update scanner failed, {e}")
            else:
                app.logger.info(f"{asset_task.taskGroupType} asset_task: {asset_task.id} "
                                f"task_id: {settings.taskId} update scanner success")

    def ws_task_log(self):
        col = ScanEvents._get_db()["scan_tasks"]
        asset_task_ids = col.distinct(
            "addition.refId",
            {"addition.taskType": {"$in": ["vul", "content", "ssl", "securityEvent"]},
             "task_config.is_periodic": True, "task_config.is_disabled": False,
             "job_config_map.site_info.cloud.http_args.area": self.old_node
             })
        app.logger.info(f"web安全监测 {len(asset_task_ids)}个asset_tasks需处理")
        return asset_task_ids

    @staticmethod
    def ws_task_run(asset_task_ids):
        asset_tasks = AssetTasks.objects.filter(
            id__in=asset_task_ids,
            taskGroupType="web_security",
            enableMonitor=True,
            triggerType="schedule"
        )
        total = asset_tasks.count()
        count = 0
        for asset_task in asset_tasks:
            count = AvailableNode.log_process(count, total)
            try:
                enable_web_sec_task(asset_task)
            except Exception as e:
                app.logger.warning(f"ws asset_task: {asset_task.id} update scanner failed, {e}")
            else:
                app.logger.info(f"ws asset_task: {asset_task.id} update scanner success")
